#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander

import tf.transformations as tfs
import numpy as np

from moveit_msgs.msg import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
from geometry_msgs.msg import PoseStamped, Pose
from std_msgs.msg import UInt8
from std_msgs.msg import Float64MultiArray
from webots_ros.msg import RecognitionObject


class MoveItIkDemo:
    def __init__(self):
        # 初始化move_group的API
        moveit_commander.roscpp_initialize(sys.argv)
        
        # 初始化ROS节点
        rospy.init_node('moveit_ik_demo')

        self.enable_flag = 0
        self.objects_info = ['can','model']  
        self.camera_to_base_R_T = None
        camera_to_chassis_base_R = None
        self.cal_R_T()

        # 初始化需要使用move group控制的机械臂中的self.arm group
        self.arm = moveit_commander.MoveGroupCommander('ls')
                
        # 获取终端link的名称
        self.end_effector_link = self.arm.get_end_effector_link()
                        
        # 设置目标位置所使用的参考坐标系
        reference_frame = 'base_link'
        self.arm.set_pose_reference_frame(reference_frame)
                
        # 当运动规划失败后,允许重新规划
        self.arm.allow_replanning(True)
        
        # 设置位置(单位：米)和姿态（单位：弧度）的允许误差
        self.arm.set_goal_position_tolerance(0.001)
        self.arm.set_goal_orientation_tolerance(0.001)
        
        rospy.Subscriber('/sign', UInt8, self.get_object_info_callback)
        rospy.Subscriber('/lsarm/kinect_color/recognition_objects', RecognitionObject, self.get_object_info_callback)
        rospy.spin()

    def get_sign_callback(self,msg_sign):
        self.enable_flag = msg_sign.data

    def get_object_info_callback(self,msg_object_info):
        if(self.enable_flag == self.objects_info.index(msg_object_info.model)):
            ob_to_camera = (np.array([[msg_object_info.position.x,msg_object_info.position.y,msg_object_info.position.z,1]])).T

            ob_to_base = np.dot(self.camera_to_base_R_T,ob_to_camera)
            print("--------------")
            print(ob_to_base)

            # self.set_pose(ob_to_base)

    def cal_R_T(self):
            camera_to_chassis_base_R = tfs.euler_matrix(-0.261795,-1.5708,0,'sxyz')
            camera_to_chassis_base_T = (np.array([[0,1.08+0.05,0,1]])).T
            for i in range(3):
                camera_to_chassis_base_R[i,3] = camera_to_chassis_base_T[i,0]

            chassis_base_to_arm_R = (tfs.euler_matrix(-1.5708,0,0)).T
            chassis_base_to_arm_T = -np.dot(chassis_base_to_arm_R,(np.array([[0,0.06,0,1]])).T)
            for i in range(3):
                chassis_base_to_arm_R[i,3] = chassis_base_to_arm_T[i,0]

            self.camera_to_base_R_T = np.dot(chassis_base_to_arm_R,camera_to_chassis_base_R)

    def set_pose(self,taget_pose):
        # target_pose = self.arm.get_current_pose()

        # 控制机械臂先回到初始化位置
        # self.arm.set_named_target('home')
        # self.arm.go()
        # rospy.sleep(5)
               
        # # 设置机械臂工作空间中的目标位姿,位置使用x、y、z坐标描述,
        # # 姿态使用四元数描述,基于base_link坐标系
        # target_pose = PoseStamped()
        # target_pose.header.frame_id = reference_frame
        # target_pose.header.stamp = rospy.Time.now()     
        target_pose.pose.position.x = 0.6
        target_pose.pose.position.y = 0.002546995731975267
        target_pose.pose.position.z = 1
        target_pose.pose.orientation.x = -0.7061563236139768
        target_pose.pose.orientation.y = -0.0049005047192389965
        target_pose.pose.orientation.z = -0.7079619781819303
        target_pose.pose.orientation.w = 0.010443616338084777
        
        #   position: 
        #     x: 0.4485708974892765
        #     y: 0.002546995731975267
        #     z: 1.039654115489455
        #   orientation: 
        #     x: -0.7061563236139768
        #     y: -0.0049005047192389965
        #     z: -0.7079619781819303
        #     w: 0.010443616338084777

        # 设置机器臂当前的状态作为运动初始状态
        self.arm.set_start_state_to_current_state()
        
        # 设置机械臂终端运动的目标位姿
        self.arm.set_pose_target(target_pose, self.end_effector_link)
        
        # 规划运动路径18.04
        # traj = self.arm.plan()
        # 规划运动路径20.04
        plan_success, traj, planning_time, error_code = self.arm.plan() 
        # 按照规划的运动路径控制机械臂运动
        self.arm.execute(traj)
        rospy.sleep(1)
         
        # # 控制机械臂回到初始化位置
        # self.arm.set_named_target('home')
        # self.arm.go()
        # 关闭并退出moveit
        moveit_commander.roscpp_shutdown()
        moveit_commander.os._exit(0)

if __name__ == "__main__":
    MoveItIkDemo()

    
    
